use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        TypedHeader,
    },
    response::IntoResponse,
    routing::get,
    Router,
};
use std::net::SocketAddr;
use tokio::net::TcpListener;
use tower_http::trace::{DefaultMakeSpan, TraceLayer};

#[tokio::main]
async fn main() {
    if std::env::var_os("RUST_LOG").is_none() {
        std::env::set_var("RUST_LOG", "example_websockets=debug,tower_http=debug")
    }
    tracing_subscriber::fmt::init();

    let listener = TcpListener::bind("localhost:3465").await.unwrap();

    let app = Router::new()
        .route("/", get(|| async { "Hello, World!" }))
        //绑定websocket路由
        .route("/sgw/ws", get(ws_handler))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::default().include_headers(true)),
        );
    let addr = SocketAddr::from(([192, 168, 128, 94], 9002));

    tracing::debug!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();

    loop {
        tokio::select! {
            _ = async {
                loop {
                    let (_socket, _) = listener.accept().await.unwrap();
                    tokio::spawn(async move { println!("1111 listener") });
                }
            } => {}
        }
    }
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    user_agent: Option<TypedHeader<headers::UserAgent>>,
) -> impl IntoResponse {
    if let Some(TypedHeader(user_agent)) = user_agent {
        println!("`{}` connected", user_agent.as_str());
    }

    ws.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    loop {
        if let Some(msg) = socket.recv().await {
            if let Ok(msg) = msg {
                println!("Client says: {:?}", msg);
            //客户端发什么，服务端就回什么（只是演示而已）
            //if socket
            //   .send(Message::Text(format!("{:?}", msg)))
            //  .await
            // .is_err()
            //{
            //   println!("client disconnected");
            //return;
            //}
            } else {
                println!("client disconnected");
                return;
            }
        }
    }
}
